<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>第四章：Chains模块 - 4.5 自定义Chain</title>
    <link rel="stylesheet" href="../shared_styles.css">
</head>
<body>
    <div class="container">
        <h2>4.5 自定义Chain</h2>
            <p>虽然Langchain提供了许多内置的Chain和Runnable，但有时我们可能需要创建具有特定逻辑或行为的自定义组件。在LCEL中，可以通过创建自定义的<code>Runnable</code>类（继承自<code>langchain_core.runnables.Runnable</code>）或使用<code>RunnableLambda</code>包装普通函数来实现。</p>
            <p><strong>示例1：一个简单的自定义Runnable，将输入文本转换为大写，并添加前缀。</strong></p>
            <pre><code class="language-python">
from typing import Any, Dict
from langchain_core.runnables import RunnableLambda
from functools import partial  # 用于向RunnableLambda传递额外参数


# 自定义处理函数
def format_text_func(inputs: Dict[str, Any], prefix: str) -> Dict[str, Any]:
    input_text = inputs.get("text_input", "")  # 从输入字典获取text_input
    processed_text = input_text.upper()
    final_text = f"{prefix}{processed_text}"
    # 返回一个包含输出的字典，也可以选择包含原始输入
    return {**inputs, "formatted_text": final_text}


# 使用RunnableLambda和partial创建自定义Runnable实例
custom_formatter_runnable = RunnableLambda(
    partial(format_text_func, prefix="Processed Output: ")
)

# 调用自定义Runnable
result_custom = custom_formatter_runnable.invoke({"text_input": "hello world"})
print(f"自定义Runnable的输入: {{'text_input': 'hello world'}}")
print(f"自定义Runnable的输出: {result_custom}")
print(f"格式化后的文本: {result_custom.get('formatted_text')}")

# 示例2: 自定义Runnable，内部组合其他Runnable (例如LLM调用)
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import Runnable, RunnableConfig  # RunnableConfig 用于类型提示
from langchain_core.messages import AIMessage  # 用于类型提示
from dotenv import load_dotenv
import os

load_dotenv()
api_key = os.environ.get("DASHSCOPE_API_KEY")
qwen_llm = ChatOpenAI(base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", api_key=api_key, model="qwen-plus")


class CustomSummarizeAndRateRunnable(Runnable):
    summarizer_runnable: Runnable  # 接收一个Runnable，例如 prompt | llm
    rater_runnable: Runnable  # 接收一个Runnable

    input_key: str = "document"  # 输入字典中期望的键

    # 添加 __init__ 方法来接收参数
    def __init__(self, summarizer_runnable: Runnable, rater_runnable: Runnable, input_key: str = "document"):
        super().__init__() # 调用父类的 __init__ 方法，虽然Runnable基类可能没有显式的__init__，但这是一个好的实践
        self.summarizer_runnable = summarizer_runnable
        self.rater_runnable = rater_runnable
        self.input_key = input_key

    def invoke(self, inputs: Dict[str, Any], config: RunnableConfig = None) -> Dict[str, Any]:
        document_content = inputs.get(self.input_key, "")

        # 第一步：调用摘要Runnable
        # summarizer_runnable 期望输入 {"text_to_summarize": ...} 并返回 AIMessage
        summary_ai_message: AIMessage = self.summarizer_runnable.invoke({"text_to_summarize": document_content})
        summary_text = summary_ai_message.content

        # 第二步：调用评价Runnable，使用上一步的摘要作为输入
        # rater_runnable 期望输入 {"content_to_rate": ...} 并返回 AIMessage
        rating_ai_message: AIMessage = self.rater_runnable.invoke({"content_to_rate": summary_text})
        rating_text = rating_ai_message.content

        return {"summary": summary_text, "rating_text": rating_text, "original_document": document_content}


# 创建子Runnable
summarize_prompt = PromptTemplate.from_template("请将以下文本总结为一句话：\n{text_to_summarize}")
summarizer = summarize_prompt | qwen_llm

rate_prompt = PromptTemplate.from_template("请评价以下内容的质量 (好/中/差)：\n{content_to_rate}")
rater = rate_prompt | qwen_llm

# 实例化自定义组合Runnable
custom_combined_runnable = CustomSummarizeAndRateRunnable(
    summarizer_runnable=summarizer,
    rater_runnable=rater
)

long_text = "Langchain是一个强大的框架，用于构建基于大型语言模型的应用程序。它提供了模块化的组件，如模型、提示、链、代理和内存，使得开发者可以轻松地将这些组件组合起来，创建复杂的AI应用。通过Langchain，可以实现问答、文本生成、摘要、翻译等多种功能。"
result_combined_custom = custom_combined_runnable.invoke({custom_combined_runnable.input_key: long_text})

print(f"\n自定义组合Runnable的输入文档长度: {len(long_text)}字")
print(f"摘要: {result_combined_custom['summary']}")
print(f"评价: {result_combined_custom['rating_text']}")

            </code></pre>
            <p>通过自定义<code>Runnable</code>或使用<code>RunnableLambda</code>，可以封装任意复杂的逻辑，并将其无缝集成到Langchain的LCEL生态系统中，实现高度的灵活性和可组合性。</p>
        <div class="navigation">
            <a href="../langchain_tutorial_outline.html">返回目录</a>
            <a href="../chapter4/4-4_router_chains.html">上一小节：Router Chains：根据输入动态选择下一个Chain (进阶)</a>
            <a href="../chapter5_agents.html">下一章：Agents模块</a>
        </div>
    </div>
</body>
</html>
